import sys
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from PySide6.QtWidgets import QWidget, QApplication, QMainWindow, QVBoxLayout, QSizePolicy
from PySide6.QtCore import Qt  # 可选：用于设置对齐方式

class MatplotlibWidget(FigureCanvas):
    """自定义Matplotlib画布部件"""
    def __init__(self, parent=None, width=8, height=4, dpi=100):
        self.fig, self.ax = plt.subplots(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)
        self.setSizePolicy(
            QSizePolicy.Policy.Expanding, 
            QSizePolicy.Policy.Expanding
        )
        self.updateGeometry()

    def clear(self):
        """清除绘图"""
        # print("> clearing plot ...")
        # self.ax.clear()
        plt.gca().cla()
        self.draw()


    def plot_test(self):
        """绘制示例图形"""
        self.ax.clear()
        
        # 生成数据
        x = np.linspace(0, 10, 100)
        y = np.sin(x)
        z = np.cos(x)
        
        # 绘制曲线
        self.ax.plot(x, y, label='sin(x)', color='blue')
        self.ax.plot(x, z, label='cos(x)', linestyle='--', color='red')
        
        # 设置图表属性
        self.ax.set_title('Matplotlib in QtCreator')
        self.ax.set_xlabel('X Axis')
        self.ax.set_ylabel('Y Axis')
        self.ax.grid(True)
        self.ax.legend()
        
        # 自动调整布局
        self.fig.tight_layout()
        self.draw()

    def plot(self, x=None, y=None):
        if x is None:
            print('x is None')
            return

        if y is None:
            print('y is None')
            return

        """绘制示例图形"""
        self.ax.clear()
        
        # 生成数据
        x = np.linspace(0, 10, 100)
        y = np.sin(x)
        z = np.cos(x)
        
        # 绘制曲线
        self.ax.plot(x, y, label='sin(x)', color='blue')
        self.ax.plot(x, z, label='cos(x)', linestyle='--', color='red')
        
        # 设置图表属性
        self.ax.set_title('Matplotlib in QtCreator')
        self.ax.set_xlabel('X Axis')
        self.ax.set_ylabel('Y Axis')
        self.ax.grid(True)
        self.ax.legend()
        
        # 自动调整布局
        self.fig.tight_layout()
        self.draw()

    def plot_gain(self, ch=None, gain=None):
        """绘制示例图形"""
        self.ax.clear()

        # 绘制曲线
        self.ax.plot(ch+1, gain, 'b--o', label='Gain')

        # 设置图表属性
        self.ax.set_title('Estimated Gain (e-/DN)')
        self.ax.set_xlabel('OS')
        self.ax.set_ylabel('Gain')
        self.ax.grid(True)
        self.ax.legend()

        # 自动调整布局
        self.fig.tight_layout()
        self.draw()

    def plot_row_overscan_gain(self, ro_gain1=None, ro_gain2=None):
        self.ax.clear()
        if ro_gain1 is None and ro_gain2 is None:
            return

        if ro_gain1 is not None:
            for i in range(ro_gain1.shape[0]):
                self.ax.plot(ro_gain1[i], '--.', ms=1, label='F1-{}'.format(i+1))

        if ro_gain2 is not None:
            for i in range(ro_gain2.shape[0]):
                self.ax.plot(ro_gain2[i], ':.', ms=1, label='F2-{}'.format(i+1))

        self.ax.grid(True)
        self.ax.legend(loc='best', frameon=False, ncol=4)

        # 自动调整布局
        self.fig.tight_layout()
        self.draw()

    def plot_flat_ratio(self, ratio18=None, ratio916=None, ylim=None):
        self.ax.clear()
        if ratio18 is not None:
            # self.ax.plot(ratio18, 'r:', alpha=0.75, label=r'OS: 1-8, RMS$\times 10^3$: {:.3f}'.format(np.std(ratio18)*1.E3))
            self.ax.plot(ratio18, 'm:', alpha=0.65, label=r'OS: 1-8')

        if ratio916 is not None:
            # self.ax.plot(ratio916, 'b:', alpha=0.75, label=r'OS: 9-16, RMS$\times 10^3$: {:.3f}'.format(np.std(ratio916)*1.E3))
            self.ax.plot(ratio916, 'g:', alpha=0.65, label=r'OS: 9-16')

        self.ax.set_xlabel('Horizontal Pixel')
        self.ax.set_ylabel('Ratio of 2 corrected flats')
        self.ax.set_xlim([-50,9266])
        if ylim is not None:
            self.ax.set_ylim(ylim)
        self.ax.grid(True)
        self.ax.legend(loc='best', frameon=True, ncol=2)

        # 自动调整布局
        self.fig.tight_layout()
        self.draw()

    def plot_ptc_gain(self, gain_result=None):
        """
        将某个CCD16个通道的增益曲线（即PTC曲线的flux/var）画出来
        """
        self.ax.clear()

        for i in range(16):
            gain = gain_result['gain_os_{}'.format(i+1)]
            self.ax.plot(gain, ':o', label='OS: {:2d}'.format(i+1))

        # 设置图表属性
        self.ax.set_title('PTC Gain (flux/var; e-/DN)')
        # self.ax.set_xlabel('OS')
        # self.ax.set_ylabel('Gain')
        self.ax.grid(True)
        self.ax.legend(ncol=4)

        # 自动调整布局
        self.fig.tight_layout()
        self.draw()

    def plot_ptc_prc_curve(self, flux=None, var=None, texp=None, plotPRC=False, saveFig=False):
        self.ax.clear()
        # self.fig, self.ax = plt.subplots(nrows=2,ncols=1,figsize=(8, 6), dpi=100)
        for ch in range(16):
            _flux = flux[f"flux_os_{ch+1}"]
            _var = var[f"var_os_{ch+1}"]

            if plotPRC:
                _texp = texp
                self.ax.plot(_texp,_flux, ':o', ms=3, label='OS: {:2d}'.format(ch+1))
            else:
                self.ax.plot(_flux, _var, ':o', ms=3, label='OS: {:2d}'.format(ch+1))

        # 设置图表属性
        if plotPRC:
            self.ax.set_title('PRC')
            self.ax.set_xlabel('ExpTime (s)')
            self.ax.set_ylabel('Flux (DN)')
        else:
            self.ax.set_title('PTC')
            self.ax.set_xlabel('Flux (DN)')
            self.ax.set_ylabel('Var (DN^2)')

        self.ax.grid(True)
        self.ax.legend(ncol=4)

        # 自动调整布局
        self.fig.tight_layout()
        self.draw()
